/*
 * Copyright 2019 The FATE Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.webank.ai.fate.serving.bean;

import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.protobuf.ByteString;
import com.webank.ai.fate.api.networking.proxy.DataTransferServiceGrpc;
import com.webank.ai.fate.api.networking.proxy.Proxy;
import com.webank.ai.fate.register.common.Constants;
import com.webank.ai.fate.register.url.URL;
import com.webank.ai.fate.serving.core.bean.Dict;
import com.webank.ai.fate.serving.core.bean.GrpcConnectionPool;
import com.webank.ai.fate.serving.core.bean.ReturnResult;
import com.webank.ai.fate.serving.core.utils.JsonUtil;
import com.webank.ai.fate.serving.core.utils.ObjectTransform;
import io.grpc.ManagedChannel;
import org.apache.commons.lang3.StringUtils;
import org.json.JSONObject;
import org.junit.BeforeClass;
import org.junit.FixMethodOrder;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.junit.runners.MethodSorters;

import java.util.List;
import java.util.concurrent.TimeUnit;

import static org.apache.logging.log4j.message.MapMessage.MapFormat.JSON;

@RunWith(JUnit4.class)
@FixMethodOrder(MethodSorters.NAME_ASCENDING)
public class UnaryCallTest {

    InferenceClient inferenceClient = new InferenceClient("localhost", 8000);

    @BeforeClass
    public static void init() {

    }

    @Test
    public void testUnaryCall() {
        try {
            Proxy.Packet.Builder packetBuilder = Proxy.Packet.newBuilder();
            packetBuilder.setBody(Proxy.Data.newBuilder()
                    .setValue(ByteString.copyFrom(JsonUtil.object2Json("{\"methodName\": \"aaaaa\"}").getBytes()))
                    .build());

            Proxy.Metadata.Builder metaDataBuilder = Proxy.Metadata.newBuilder();
            Proxy.Topic.Builder topicBuilder = Proxy.Topic.newBuilder();

            metaDataBuilder.setSrc(
                    topicBuilder.setPartyId("9999").
                            setRole("serving")
                            .setName(Dict.PARTNER_PARTY_NAME)
                            .build());
            metaDataBuilder.setDst(
                    topicBuilder.setPartyId("9999")
                            .setRole("serving")
                            .setName(Dict.PARTY_NAME)
                            .build());
            metaDataBuilder.setCommand(Proxy.Command.newBuilder().setName("federatedInference").build());
            metaDataBuilder.setConf(Proxy.Conf.newBuilder().setOverallTimeout(60 * 1000));
            packetBuilder.setHeader(metaDataBuilder.build());
            /*Proxy.AuthInfo.Builder authBuilder = Proxy.AuthInfo.newBuilder();
            if (context.getCaseId() != null) {
                authBuilder.setNonce(context.getCaseId());
            }
            if (version != null) {
                authBuilder.setVersion(version);
            }
            if (context.getServiceId() != null) {
                authBuilder.setServiceId(context.getServiceId());
            }
            if (context.getApplyId() != null) {
                authBuilder.setApplyId(context.getApplyId());
            }
            packetBuilder.setAuth(authBuilder.build());*/

            GrpcConnectionPool grpcConnectionPool = GrpcConnectionPool.getPool();
            ManagedChannel channel1 = grpcConnectionPool.getManagedChannel("127.0.0.1", 8869);

            DataTransferServiceGrpc.DataTransferServiceFutureStub stub1 = DataTransferServiceGrpc.newFutureStub(channel1);
            ListenableFuture<Proxy.Packet> future = stub1.unaryCall(packetBuilder.build());

            if (future != null) {
                Proxy.Packet packet = future.get(3000, TimeUnit.MILLISECONDS);
                ReturnResult remoteResult = (ReturnResult) ObjectTransform.json2Bean(packet.getBody().getValue().toStringUtf8(), ReturnResult.class);
                System.out.println(JsonUtil.object2Json(remoteResult));
                System.out.println("========================");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}